#include <mcp/server/prompt_manager.h>
#include <mcp/json/value.h>
#include <mcp/error.h>
#include <mcp/types.h>
#include <iostream>
#include <sstream>
#include <algorithm>
#include <memory>

namespace mcp {
namespace server {

// PromptManager implementation
struct PromptManager::Impl {
    std::unordered_map<std::string, PromptRegistration> prompts;
    std::vector<std::string> prompt_order; // For maintaining insertion order

    // Thread safety
    mutable std::mutex mutex;
};

PromptManager::PromptManager() : impl_(std::make_unique<Impl>()) {}

PromptManager::~PromptManager() = default;

bool PromptManager::register_prompt(PromptRegistration registration) {
    std::lock_guard<std::mutex> lock(impl_->mutex);

    if (impl_->prompts.contains(registration.name)) {
        return false; // Prompt already exists
    }

    impl_->prompts[registration.name] = std::move(registration);
    impl_->prompt_order.push_back(registration.name);
    return true;
}

bool PromptManager::unregister_prompt(const std::string& name) {
    std::lock_guard<std::mutex> lock(impl_->mutex);

    auto it = impl_->prompts.find(name);
    if (it == impl_->prompts.end()) {
        return false; // Prompt not found
    }

    impl_->prompts.erase(it);
    impl_->prompt_order.erase(
        std::remove(impl_->prompt_order.begin(), impl_->prompt_order.end(), name),
        impl_->prompt_order.end()
    );
    return true;
}

bool PromptManager::has_prompt(const std::string& name) const {
    std::lock_guard<std::mutex> lock(impl_->mutex);
    return impl_->prompts.contains(name);
}

std::vector<types::Prompt> PromptManager::list_prompts() const {
    std::lock_guard<std::mutex> lock(impl_->mutex);

    std::vector<types::Prompt> result;
    result.reserve(impl_->prompt_order.size());

    for (const auto& name : impl_->prompt_order) {
        auto it = impl_->prompts.find(name);
        if (it != impl_->prompts.end()) {
            const auto& registration = it->second;
            std::vector<types::PromptArgument> prompt_args =
                build_prompt_arguments(registration.arguments);

            result.push_back(types::Prompt{
                .name = registration.name,
                .description = registration.description,
                .arguments = std::move(prompt_args)
            });
        }
    }

    return result;
}

std::optional<types::Prompt> PromptManager::get_prompt(const std::string& name) const {
    std::lock_guard<std::mutex> lock(impl_->mutex);

    auto it = impl_->prompts.find(name);
    if (it == impl_->prompts.end()) {
        return std::nullopt;
    }

    const auto& registration = it->second;
    std::vector<types::PromptArgument> prompt_args =
        build_prompt_arguments(registration.arguments);

    return types::Prompt{
        .name = registration.name,
        .description = registration.description,
        .arguments = std::move(prompt_args)
    };
}

std::optional<std::vector<PromptArgumentDefinition>> PromptManager::get_arguments(const std::string& name) const {
    std::lock_guard<std::mutex> lock(impl_->mutex);

    auto it = impl_->prompts.find(name);
    if (it == impl_->prompts.end()) {
        return std::nullopt;
    }

    return it->second.arguments;
}

Result<std::vector<types::PromptMessage>> PromptManager::call_prompt(const std::string& name,
                                                                      const json::Value& arguments,
                                                                      Context<>& context) {
    std::lock_guard<std::mutex> lock(impl_->mutex);

    auto it = impl_->prompts.find(name);
    if (it == impl_->prompts.end()) {
        return Error::method_not_found(std::string("Prompt not found: ") + name);
    }

    const auto& registration = it->second;

    try {
        // Validate arguments against schemas
        auto validation = validate_arguments(name, arguments);
        if (!validation.valid) {
            std::string error_msg = "Invalid arguments for prompt '" + name + "':\n";
            for (const auto& error : validation.errors) {
                error_msg += "  - " + error + "\n";
            }
            return Error::invalid_params(error_msg);
        }

        // Check for required arguments
        for (const auto& arg_def : registration.arguments) {
            if (arg_def.required) {
                if (!arguments.contains(arg_def.name)) {
                    return Error::invalid_params(
                        std::string("Missing required argument: ") + arg_def.name
                    );
                }
            }
        }

        // Call the prompt handler
        context.info(std::format("Calling prompt '{}' with {} arguments", name, arguments.size()));
        auto result = registration.handler(arguments, context);
        if (!result) {
            return result.error();
        }

        // Validate the result
        if (result->empty()) {
            context.warning(std::format("Prompt '{}' returned empty messages", name));
        }

        // Validate message structure
        for (size_t i = 0; i < result->size(); ++i) {
            const auto& message = (*result)[i];
            if (!message.content) {
                context.error(std::format("Prompt '{}' returned message {} with null content", name, i));
                return Error::internal_error("Prompt returned message with null content");
            }
        }

        return result;

    } catch (const std::exception& e) {
        return Error::internal_error(
            std::string("Exception in prompt '") + name + "': " + e.what()
        );
    } catch (...) {
        return Error::internal_error(
            std::string("Unknown exception in prompt '") + name + "'"
        );
    }
}

json::ValidationResult PromptManager::validate_arguments(const std::string& name,
                                                         const json::Value& arguments) const {
    std::lock_guard<std::mutex> lock(impl_->mutex);

    auto it = impl_->prompts.find(name);
    if (it == impl_->prompts.end()) {
        return json::ValidationResult{false, {"Prompt not found: " + name}};
    }

    const auto& registration = it->second;
    json::ValidationResult result{true, {}};

    // Validate each argument against its schema
    for (const auto& arg_def : registration.arguments) {
        if (arguments.contains(arg_def.name)) {
            if (arg_def.schema) {
                auto arg_validation = arg_def.schema->validate(arguments[arg_def.name]);
                if (!arg_validation.valid) {
                    result.valid = false;
                    result.errors.insert(result.errors.end(),
                                        arg_validation.errors.begin(),
                                        arg_validation.errors.end());
                }
            }
        } else if (arg_def.required) {
            result.valid = false;
            result.errors.push_back("Missing required argument: " + arg_def.name);
        }
    }

    // Check for unexpected arguments
    if (arguments.is_object()) {
        for (const auto& [key, value] : arguments.as_object()) {
            bool expected = false;
            for (const auto& arg_def : registration.arguments) {
                if (arg_def.name == key) {
                    expected = true;
                    break;
                }
            }
            if (!expected) {
                result.valid = false;
                result.errors.push_back("Unexpected argument: " + key);
            }
        }
    }

    return result;
}

std::vector<types::PromptArgument> PromptManager::build_prompt_arguments(
    const std::vector<PromptArgumentDefinition>& arguments) {

    std::vector<types::PromptArgument> result;
    result.reserve(arguments.size());

    for (const auto& arg_def : arguments) {
        result.push_back(types::PromptArgument{
            .name = arg_def.name,
            .description = arg_def.description,
            .required = arg_def.required
        });
    }

    return result;
}

void PromptManager::clear() {
    std::lock_guard<std::mutex> lock(impl_->mutex);
    impl_->prompts.clear();
    impl_->prompt_order.clear();
}

// Prompt helper functions
types::PromptMessage create_user_message(const std::string& text) {
    return types::PromptMessage{
        .role = types::PromptRole::User,
        .content = std::make_shared<types::TextContent>(text)
    };
}

types::PromptMessage create_assistant_message(const std::string& text) {
    return types::PromptMessage{
        .role = types::PromptRole::Assistant,
        .content = std::make_shared<types::TextContent>(text)
    };
}

types::PromptMessage create_system_message(const std::string& text) {
    return types::PromptMessage{
        .role = types::PromptRole::System,
        .content = std::make_shared<types::TextContent>(text)
    };
}

std::vector<types::PromptMessage> create_conversation(
    const std::vector<std::pair<types::PromptRole, std::string>>& messages) {

    std::vector<types::PromptMessage> result;
    result.reserve(messages.size());

    for (const auto& [role, text] : messages) {
        result.push_back(types::PromptMessage{
            .role = role,
            .content = std::make_shared<types::TextContent>(text)
        });
    }

    return result;
}

// Prompt validation helpers
bool validate_prompt_name(const std::string& name) {
    if (name.empty()) {
        return false;
    }

    // Prompt names should be valid identifiers
    for (char c : name) {
        if (!(std::isalnum(c) || c == '_' || c == '-' || c == '.')) {
            return false;
        }
    }

    // Cannot start with a number or special character
    if (std::isdigit(name[0]) || name[0] == '-' || name[0] == '.') {
        return false;
    }

    return true;
}

std::string sanitize_prompt_name(const std::string& name) {
    std::string result;
    result.reserve(name.size());

    for (char c : name) {
        if (std::isalnum(c) || c == '_' || c == '-' || c == '.') {
            result += c;
        } else {
            result += '_';
        }
    }

    // Ensure it doesn't start with invalid characters
    while (!result.empty() && (std::isdigit(result[0]) || result[0] == '-' || result[0] == '.')) {
        result.erase(0, 1);
    }

    return result.empty() ? "unnamed_prompt" : result;
}

// Prompt argument helpers
PromptArgumentDefinition create_string_argument(
    const std::string& name,
    bool required,
    const std::optional<std::string>& description) {

    return PromptArgumentDefinition{
        .name = name,
        .description = description,
        .required = required,
        .schema = json::string_schema()
    };
}

PromptArgumentDefinition create_number_argument(
    const std::string& name,
    bool required,
    const std::optional<std::string>& description) {

    return PromptArgumentDefinition{
        .name = name,
        .description = description,
        .required = required,
        .schema = json::number_schema()
    };
}

PromptArgumentDefinition create_boolean_argument(
    const std::string& name,
    bool required,
    const std::optional<std::string>& description) {

    return PromptArgumentDefinition{
        .name = name,
        .description = description,
        .required = required,
        .schema = json::boolean_schema()
    };
}

PromptArgumentDefinition create_object_argument(
    const std::string& name,
    bool required,
    const std::optional<std::string>& description,
    const json::Schema& schema) {

    return PromptArgumentDefinition{
        .name = name,
        .description = description,
        .required = required,
        .schema = schema
    };
}

PromptArgumentDefinition create_enum_argument(
    const std::string& name,
    bool required,
    const std::vector<std::string>& enum_values,
    const std::optional<std::string>& description) {

    json::Schema schema = json::string_schema();
    for (const auto& value : enum_values) {
        schema = schema.enum_values({json::Value(value)});
    }

    return PromptArgumentDefinition{
        .name = name,
        .description = description,
        .required = required,
        .schema = std::move(schema)
    };
}

// Prompt execution context
class PromptExecutionContext {
public:
    explicit PromptExecutionContext(const std::string& prompt_name, Context<>& context)
        : prompt_name_(prompt_name), context_(context) {
        start_time_ = std::chrono::steady_clock::now();
    }

    ~PromptExecutionContext() {
        auto end_time = std::chrono::steady_clock::now();
        auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(end_time - start_time_);
        context_.info(std::format("Prompt '{}' completed in {}ms", prompt_name_, duration.count()));
    }

    void log_generation_start(size_t message_count = 1) {
        context_.info(std::format("Generating {} message(s) for prompt '{}'", message_count, prompt_name_));
    }

    void log_generation_complete(const std::vector<types::PromptMessage>& messages) {
        context_.info(std::format("Prompt '{}' generated {} messages", prompt_name_, messages.size()));
    }

private:
    std::string prompt_name_;
    Context<>& context_;
    std::chrono::steady_clock::time_point start_time_;
};

// Prompt result builders
class PromptResultBuilder {
public:
    explicit PromptResultBuilder(const std::string& prompt_name) : prompt_name_(prompt_name) {}

    PromptResultBuilder& add_user_message(const std::string& text) {
        messages_.push_back(create_user_message(text));
        return *this;
    }

    PromptResultBuilder& add_assistant_message(const std::string& text) {
        messages_.push_back(create_assistant_message(text));
        return *this;
    }

    PromptResultBuilder& add_system_message(const std::string& text) {
        messages_.push_back(create_system_message(text));
        return *this;
    }

    PromptResultBuilder& add_message(types::PromptRole role, const std::string& text) {
        messages_.push_back(types::PromptMessage{
            .role = role,
            .content = std::make_shared<types::TextContent>(text)
        });
        return *this;
    }

    std::vector<types::PromptMessage> build() && {
        return std::move(messages_);
    }

private:
    std::string prompt_name_;
    std::vector<types::PromptMessage> messages_;
};

// Template prompt generation
std::vector<types::PromptMessage> create_template_prompt(
    const std::string& system_template,
    const std::string& user_template,
    const std::unordered_map<std::string, std::string>& variables) {

    std::vector<types::PromptMessage> messages;

    // Add system message if template is provided
    if (!system_template.empty()) {
        std::string system_content = system_template;
        for (const auto& [key, value] : variables) {
            std::string placeholder = "{" + key + "}";
            size_t pos = system_content.find(placeholder);
            while (pos != std::string::npos) {
                system_content.replace(pos, placeholder.length(), value);
                pos = system_content.find(placeholder, pos + value.length());
            }
        }
        messages.push_back(create_system_message(system_content));
    }

    // Add user message
    std::string user_content = user_template;
    for (const auto& [key, value] : variables) {
        std::string placeholder = "{" + key + "}";
        size_t pos = user_content.find(placeholder);
        while (pos != std::string::npos) {
            user_content.replace(pos, placeholder.length(), value);
            pos = user_content.find(placeholder, pos + value.length());
        }
    }
    messages.push_back(create_user_message(user_content));

    return messages;
}

// Conversation prompt helpers
std::vector<types::PromptMessage> create_qa_prompt(
    const std::string& question,
    const std::string& context = "",
    const std::string& instructions = "") {

    std::vector<types::PromptMessage> messages;

    // Add system instructions if provided
    if (!instructions.empty()) {
        messages.push_back(create_system_message(instructions));
    }

    // Add context if provided
    if (!context.empty()) {
        messages.push_back(create_system_message("Context: " + context));
    }

    // Add the question
    messages.push_back(create_user_message(question));

    return messages;
}

std::vector<types::PromptMessage> create_code_review_prompt(
    const std::string& code,
    const std::string& language = "",
    const std::string& focus = "") {

    std::string system_message = "You are a code reviewer. ";
    if (!language.empty()) {
        system_message += "The code is written in " + language + ". ";
    }
    if (!focus.empty()) {
        system_message += "Focus your review on: " + focus + ". ";
    }
    system_message += "Provide constructive feedback on code quality, bugs, performance, and best practices.";

    return {
        create_system_message(system_message),
        create_user_message("Please review this code:\n\n" + code)
    };
}

std::vector<types::PromptMessage> create_summarization_prompt(
    const std::string& text,
    const std::string& summary_type = "brief",
    size_t target_length = 100) {

    std::string instructions = std::format(
        "Summarize the following text in {} words or less. "
        "Focus on the key points and main ideas. "
        "The summary should be {} and informative.",
        target_length,
        summary_type
    );

    return {
        create_system_message(instructions),
        create_user_message(text)
    };
}

} // namespace server
} // namespace mcp